// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::collections::BTreeMap;
use std::ffi::c_void;
use std::ptr;

use anyhow::bail;
use dashmap::DashMap;
use once_cell::sync::Lazy;
use parking_lot::Mutex;
use tracing::{debug, error};

#[derive(Debug)]
struct VHandleIntervalMap {
    map: BTreeMap<u64, u64>,
}

impl VHandleIntervalMap {
    fn new() -> Self {
        VHandleIntervalMap {
            map: BTreeMap::new(),
        }
    }

    fn insert(&mut self, start: u64, end: u64) {
        self.map.insert(start, end);
    }

    fn query(&self, addr: u64) -> Option<(u64, u64)> {
        let entry = self.map.range(..=addr).next_back();
        if let Some((&start, &end)) = entry {
            if addr < end || addr == end && start == end {
                return Some((start, addr - start));
            }
        }

        None
    }
}

static HANDLE_MAP: Lazy<DashMap<u64, u64>> = Lazy::new(DashMap::new);
static VHANDLE_INTERVALS: Lazy<Mutex<VHandleIntervalMap>> =
    Lazy::new(|| Mutex::new(VHandleIntervalMap::new()));
static REQ_ID_VHANDLE_MAP: Lazy<DashMap<u64, u64>> = Lazy::new(DashMap::new);

pub fn is_vhandle_valid(vh: *mut c_void) -> bool {
    matches!(VHANDLE_INTERVALS.lock().query(vh as u64), Some((_, _)))
}

pub fn handle_insert(vh: *mut c_void, ph: *mut c_void, size: usize) -> anyhow::Result<()> {
    if ph.is_null() {
        error!("[Virt] Failed to insert: ph is null");
        bail!("Failed to insert: ph is null");
    }

    let key = vh as u64;
    let value = ph as u64;
    if HANDLE_MAP.insert(key, value).is_some() {
        debug!("[Virt] Updated mapping vh {:p} -> ph {:p}", vh, ph);
    } else {
        let start = key;
        let end = key.wrapping_add(size as u64);
        let mut intervals = VHANDLE_INTERVALS.lock();
        intervals.insert(start, end);
        debug!("[Virt] Inserted mapping vh {:p} -> ph {:p}", vh, ph);
        debug!(
            "[Virt] Inserted memory interval ({:>#x}, {:>#x})",
            start, end
        );
    }

    Ok(())
}

pub fn handle_map(vh: *mut c_void) -> anyhow::Result<*mut c_void> {
    if vh.is_null() {
        debug!("[Virt] Mapping NULL handle, may be default API param");
        return Ok(ptr::null_mut());
    }

    if let Some((base, offset)) = VHANDLE_INTERVALS.lock().query(vh as u64) {
        debug!(
            "[Virt] Found base for vh {:p}: (base {:>#x}, offset {:>#x})",
            vh, base, offset
        );
        if let Some(v) = HANDLE_MAP.get(&base) {
            let ph = (v.value() + offset) as *mut c_void;
            debug!("[Virt] Mapped vh {:p} -> ph {:p}", vh, ph);
            Ok(ph)
        } else {
            error!("[Virt] Error mapping vh: key {:>#x} not found", base);
            bail!("Error mapping vh: key {:>#x} not found", base);
        }
    } else {
        error!("[Virt] Error mapping vh: {:p} not in any interval", vh);
        bail!("Error mapping vh: {:p} not in any interval", vh);
    }
}

pub fn req_id_vhandle_insert(request_id: u64, vh: *mut c_void) -> anyhow::Result<()> {
    let value = vh as u64;
    REQ_ID_VHANDLE_MAP.insert(request_id, value);
    debug!("[Virt] Inserted request_id {} -> vh {:p}", request_id, vh);

    Ok(())
}

pub fn req_id_vhandle_map(request_id: u64) -> anyhow::Result<*mut c_void> {
    if let Some(v) = REQ_ID_VHANDLE_MAP.get(&request_id) {
        let vh = *v.value() as *mut c_void;
        debug!("[Virt] Mapped request_id {} -> vh {:p}", request_id, vh);

        Ok(vh)
    } else {
        error!("[Virt] Error mapping request_id {}: not found", request_id);
        bail!("Error mapping request_id {}: not found", request_id);
    }
}
